__author__ = 'sdjkx'

import os
import shutil
from libs.ustr import *

try:
  import xml.etree.cElementTree as ET
except ImportError:
  import xml.etree.ElementTree as ET
from PyQt5.QtWidgets import QMessageBox, QFileDialog

class EasyDLFile:

    def __init__(self,labelImg):
        self.labelImg = labelImg

    def create(self, imgPath, xmlPath,label_cn_en):
        if imgPath is not None:
            if imgPath.find("train") < 0:
                QMessageBox.information(self.labelImg, "错误","图片目录必须是【train】目录")
                return

            ret = QMessageBox.question(self.labelImg, '提示', "确定要将目录【"+imgPath+"】里图片和XML文件生成到EasyDL吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if ret == QMessageBox.Yes:
                self.scanFileToEasyDl(imgPath,xmlPath,label_cn_en)
                return
        else:
            imgPath = '.'

        dirpath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择XML文件路径", imgPath,  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if dirpath is not None and len(dirpath) > 1:
            if dirpath.find("train") < 0:
                QMessageBox.information(self.labelImg, "错误","图片目录必须是【train】目录")
                return

            ret = QMessageBox.question(self.labelImg, '提示',  "确定要将目录【"+dirpath+"】里图片和XML文件生成到EasyDL吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if ret == QMessageBox.Yes:
                xmlPath = dirpath.replace("train", "Annotations")
                self.scanFileToEasyDl(dirpath,xmlPath,label_cn_en)
                return

    def getBasePath(self):
        basePath = "F:/Phone/outdoor/easyDL/"
        ret = QMessageBox.question(self.labelImg, '提示',  "是否使用目录【"+basePath+"】保存EasyDL文件？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if ret == QMessageBox.No:
            basePath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                               "选择保存EasyDL文件路径", basePath,  QFileDialog.ShowDirsOnly
                                               | QFileDialog.DontResolveSymlinks))
            if basePath is None:
                return ""

        return basePath

    def scanFileToEasyDl(self, imgPath, xmlPath,label_cn_en):
        self.label_cn_en = label_cn_en
        basePath = self.getBasePath()
        if basePath == "":
            return

        if not os.path.exists(basePath):
            os.mkdir(basePath)

        if basePath == imgPath:
            QMessageBox.information(self.labelImg, "错误","EasyDL文件不可以生成到图片目录")
            return
        if basePath == xmlPath:
            QMessageBox.information(self.labelImg, "错误","EasyDL文件不可以生成到XML目录")
            return

        QMessageBox.information(self.labelImg, "提示","开始生成EasyDL文件")

        img_list = os.listdir(imgPath)

        for imgFileName in img_list:
            # print("xml_n=%s" % xml_n)
            img_path = os.path.join(imgPath, imgFileName)
            detImg_path =  os.path.join(basePath , imgFileName)

            xml_path = xmlPath + "/" + imgFileName[0:-3] + "xml"
            if not os.path.exists(xml_path):
                print("XML文件不存在" + xml_path)
                continue

            if not os.path.exists(detImg_path):
                shutil.copy(img_path, detImg_path)


            tree = ET.parse(xml_path)     #打开xml文档
            root = tree.getroot()         #获得root节点
            # print ("*"*10)

            box_count = 0
            all_box_str = ""

            for object in root.findall('object'): #找到root节点下的所有object节点
                label_name = object.find('name').text   #子节点下节点name的值
                # if self.label_cn_en.__contains__(label_name):
                if label_name in self.label_cn_en:
                    en_name = self.label_cn_en[label_name]
                else:
                    return ""

                box_count += 1

                bndbox = object.find('bndbox')      #子节点下属性bndbox的值

                xmin = float(bndbox.find('xmin').text)
                ymin = float(bndbox.find('ymin').text)
                xmax = float(bndbox.find('xmax').text)
                ymax = float(bndbox.find('ymax').text)
                # {
                #     "labels": [
                #         {
                #             "name": "car",
                #             "x1": 278.7673,
                #             "y1": 38.76027,
                #             "x2": 328.51887,
                #             "y2": 101.77642
                #         },
                #         {
                #             "name": "car",
                #             "x1": 305.5889,
                #             "y1": 48.56245,
                #             "x2": 351.31532000000007,
                #             "y2": 112.22009
                #         },
                #         ...
                # }

                # 图片路径\t	{"value":"bolt","coordinate":[[769.459,241.819],[947.546,506.167]]}\t{...}... 注意，每个字段值之间用\t分割
                # box_str = '{{\"value\":\"Leconte\",\"coordinate\":[[{:.3f},{:.3f}],[{:.3f},{:.3f}]]}}\t'.\
                #                                                 format(xmin,  ymin,  xmax, ymax)
                box_str = '{\"name\": \"'+en_name+'\","x1": '+str(xmin)+',"y1": '+str(ymin)+',"x2": '+str(xmax)+',"y2": '+str(ymax)+'},'
                all_box_str += box_str

            if box_count==0:
                continue

            all_box_str = '{"labels": [' + all_box_str[0:-1] + "]}"

            out_txt_path_train = basePath +  imgFileName[0:-3] + "json"
            txt_json = open(out_txt_path_train, 'w')
            txt_json.write(all_box_str)
            txt_json.close()

        QMessageBox.information(self.labelImg, "","数据处理完成")
        return
